import pytesseract
from PIL import Image
import re
import time
import cv2
import os

class PicTreat:
    def __init__(self, driver, css):
        self.driver = driver
        self.css = css

    def pic_to_num(self, pic):
        text = pytesseract.image_to_string(Image.open(pic), lang="eng")
        text = re.search('\S+', text)[0]
        return text

    def get_snap(self):  # 对目标网页进行截屏。这里截的是全屏
        self.driver.save_screenshot('full_snap.png')
        page_snap_obj = Image.open('full_snap.png')
        return page_snap_obj

    def get_image(self, name):  # 对验证码所在位置进行定位，然后截取验证码图片
        img = self.driver.find_element_by_css_selector(self.css)
        time.sleep(2)
        location = img.location
        print(location)
        size = img.size
        left = location['x']
        top = location['y']
        right = left + size['width']
        bottom = top + size['height']

        page_snap_obj = self.get_snap()
        image_obj = page_snap_obj.crop((left, top, right, bottom))
        # image_obj.show()
        image_obj.save(name)

    def get_num(self):
        TimeTuple = time.localtime(time.time())  # 获取当前的时间返回一个时间元组
        fmt = '%Y%m%d%H%M%S'
        now = time.strftime(fmt, TimeTuple)
        self.get_image('num_pic\\{}.png'.format(now))
        num = self.pic_to_num('num_pic\\{}.png'.format(now))
        return num

# if __name__ == '__main__':
#     pic = Pic_treat()